"""
训练LightGBM模型并导出为ONNX格式
该脚本从features_daily表中读取特征数据，训练两个LightGBM模型（分别预测开盘和收盘涨跌），
然后将训练好的模型转换为ONNX格式以便后续推理使用。
"""

import pandas as pd                           # 数据处理库
import lightgbm as lgb                        # LightGBM机器学习库
from sqlalchemy import create_engine          # 数据库连接库
from config import DB_CONFIG, FEATURE_ORDER, MODEL_DIR  # 导入配置信息
from transformer import save_transformer     # 特征转换器保存函数
import os                                    # 系统操作库
import onnxmltools                           # ONNX模型转换工具
from onnxmltools.convert.common.data_types import FloatTensorType  # ONNX数据类型

# 创建模型存储目录
os.makedirs(MODEL_DIR, exist_ok=True)

def get_engine():
    """
    创建数据库连接引擎
    Returns:
        sqlalchemy.engine.Engine: PostgreSQL数据库连接引擎
    """
    uri = f"postgresql+psycopg2://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
    return create_engine(uri)

def main():
    """
    主训练函数
    负责数据加载、模型训练和ONNX模型导出的完整流程
    """
    # 创建数据库连接
    engine = get_engine()
    
    # 从数据库加载特征数据，只选择标签非空的记录
    df = pd.read_sql("SELECT * FROM features_daily WHERE label_next_open_up IS NOT NULL", engine)
    
    # 删除特征列和标签列中的空值行，确保数据质量
    df = df.dropna(subset=FEATURE_ORDER + ["label_next_open_up", "label_next_close_up"])

    # 保存特征转换器（用于推理时数据预处理）
    save_transformer(df)

    # 准备特征数据和标签数据
    X = df[FEATURE_ORDER].astype(float)           # 特征矩阵，转换为浮点型
    y_open = df["label_next_open_up"].astype(int)     # 开盘涨跌标签（0=跌，1=涨）
    y_close = df["label_next_close_up"].astype(int)   # 收盘涨跌标签（0=跌，1=涨）

    # 划分训练集和测试集（80%训练，20%测试）
    split = int(len(df) * 0.8)
    X_train, X_test = X.iloc[:split], X.iloc[split:]              # 特征数据划分
    y_train_open, y_test_open = y_open.iloc[:split], y_open.iloc[split:]    # 开盘标签划分
    y_train_close, y_test_close = y_close.iloc[:split], y_close.iloc[split:]  # 收盘标签划分

    # 训练开盘涨跌预测模型
    # 使用LightGBM分类器，设置200棵树，学习率0.05，叶子节点数31，随机种子42保证结果可重现
    model_open = lgb.LGBMClassifier(n_estimators=200, learning_rate=0.05, num_leaves=31, random_state=42)
    model_open.fit(X_train, y_train_open)  # 拟合模型

    # 训练收盘涨跌预测模型
    # 使用相同参数训练收盘价预测模型
    model_close = lgb.LGBMClassifier(n_estimators=200, learning_rate=0.05, num_leaves=31, random_state=42)
    model_close.fit(X_train, y_train_close)  # 拟合模型

    # 定义ONNX模型输入类型（浮点型张量，第一维None表示批大小可变）
    initial_type = [("float_input", FloatTensorType([None, len(FEATURE_ORDER)]))]

    # 将开盘预测模型转换为ONNX格式并保存
    onnx_open = onnxmltools.convert_lightgbm(model_open.booster_, initial_types=initial_type, target_opset=15)
    with open(f"{MODEL_DIR}/open_up.onnx", "wb") as f:
        f.write(onnx_open.SerializeToString())  # 写入文件

    # 将收盘预测模型转换为ONNX格式并保存
    onnx_close = onnxmltools.convert_lightgbm(model_close.booster_, initial_types=initial_type, target_opset=15)
    with open(f"{MODEL_DIR}/close_up.onnx", "wb") as f:
        f.write(onnx_close.SerializeToString())  # 写入文件

    print("✅ 模型训练完成并导出 ONNX。")

if __name__ == "__main__":
    main()